import numpy as np
import tensorflow as tf
import os
import re
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras import utils
from sklearn.preprocessing import MinMaxScaler, StandardScaler

def load_data(data_path, input_shape, labels_path=None):
	"""
	加载数据。默认支持 .npy 文件且当 labels_path 为 None 时假设数据最后一列为标签。
	input_shape 用于重塑当原数据以展平形式存储时恢复 (time_steps, channels)。
	"""
	data = np.load(data_path)
	if labels_path:
		labels = np.load(labels_path)
	else:
		labels = data[:, -1]
		data = data[:, :-1].reshape(-1, input_shape[0], input_shape[1])
	return data, labels

def load_dataset_from_files(root_dir, motion_to_label, DEF_FILE_FORMAT=".txt", DEF_FILE_MAX=999, DEF_USE_COLS=(0, 1, 2, 3, 4, 5), DEF_N_ROWS=None):
	"""
	从文件夹加载多个数据文件,每个文件包含一个动作样本。
	
	参数:
	- root_dir: 数据文件所在目录
	- motion_to_label: 动作名称到标签的映射字典
	- DEF_FILE_FORMAT: 文件格式后缀
	- DEF_FILE_MAX: 文件名中允许的最大数字
	- DEF_USE_COLS: 使用的数据列
	- DEF_N_ROWS: 读取的最大行数
	
	返回:
	- file_list: 数据列表
	- labels: 标签列表
	"""
	# 检查目录是否存在
	if not os.path.exists(root_dir):
		raise FileNotFoundError(f"目录不存在: {root_dir}")
	
	file_list = []
	labels = []
	
	# 检查motion_to_label是否为空
	if not motion_to_label:
		raise ValueError("motion_to_label字典为空,无法映射动作名称到标签。")
	
	# 获取目录中的所有文件
	try:
		files = os.listdir(root_dir)
	except Exception as e:
		raise IOError(f"无法读取目录 {root_dir}: {e}")
	
	# 如果目录为空,发出警告
	if len(files) == 0:
		print(f"警告: 目录 {root_dir} 中没有文件。")
		return file_list, labels
	
	for filename in files:
		if filename.endswith(DEF_FILE_FORMAT):
			match = re.match(rf'^([\w]+)_([\d]+){DEF_FILE_FORMAT}$', filename)
			if match:
				motion_name = match.group(1)
				number_str = match.group(2)
				number = int(number_str)
				if 0 <= number <= DEF_FILE_MAX:
					if motion_name in motion_to_label:
						file_path = os.path.join(root_dir, filename)
						try:
							# 使用max_rows参数限制读取的行数
							data = np.loadtxt(file_path, delimiter=' ', usecols=DEF_USE_COLS, max_rows=DEF_N_ROWS)
							# 检查数据是否为空
							if data.size == 0:
								print(f"警告: 文件 {filename} 中没有数据。")
								continue
							# 确保数据是二维数组
							if len(data.shape) == 1:
								data = data.reshape(-1, len(DEF_USE_COLS) if DEF_USE_COLS else data.shape[0])
							file_list.append(data)
							labels.append(motion_to_label[motion_name])
						except Exception as e:
							print(f"加载文件 {filename} 时出错: {e}")
					else:
						print(f"Motion name not recognized: {filename}")
				else:
					print(f"Number out of range: {filename}")
			else:
				print(f"Invalid file name format: {filename}")
	
	# 检查是否加载了任何数据
	if len(file_list) == 0:
		print(f"警告: 没有从目录 {root_dir} 中加载任何有效数据。")
	
	return file_list, labels

def preprocess_dataset(file_list, labels, num_classes, input_shape=None):
	"""
	预处理数据集,包括填充序列、转换为one-hot编码、分割数据集。
	
	参数:
	- file_list: 数据列表
	- labels: 标签列表
	- num_classes: 类别数量
	- input_shape: 期望的输入形状 (time_steps, channels),如果为None则使用最长序列的长度
	
	返回:
	- x_train: 训练数据
	- y_train: 训练标签
	- x_test: 测试数据
	- y_test: 测试标签
	"""
	# 检查数据列表是否为空
	if len(file_list) == 0:
		raise ValueError("数据列表为空,无法进行预处理。")
	
	# 检查每个数据元素的形状
	for i, data in enumerate(file_list):
		if len(data.shape) != 2:
			print(f"警告: 数据索引 {i} 的形状不是2维,而是 {data.shape},尝试修复...")
			# 如果数据是一维的,尝试重塑为二维
			if len(data.shape) == 1:
				file_list[i] = data.reshape(-1, 1)
			else:
				raise ValueError(f"数据索引 {i} 的形状 {data.shape} 无法处理。")
	
	# 数据预处理,例如填充序列以达到统一长度
	max_len = max([len(x) for x in file_list])  # 找到最长序列的长度
	print(f"Max length of sequences: {max_len}")  # 打印max_len的值
	
	# 如果指定了input_shape,使用指定的time_steps
	if input_shape is not None:
		max_len = input_shape[0]
	
	# 确保所有数据都有相同的通道数
	channels = file_list[0].shape[1] if len(file_list) > 0 else 1
	for i, data in enumerate(file_list):
		if data.shape[1] != channels:
			print(f"警告: 数据索引 {i} 的通道数 {data.shape[1]} 与期望的 {channels} 不匹配,尝试修复...")
			# 截断或填充通道
			if data.shape[1] > channels:
				file_list[i] = data[:, :channels]
			else:
				padding = np.zeros((data.shape[0], channels - data.shape[1]))
				file_list[i] = np.hstack((data, padding))
	
	file_list_padded = pad_sequences(file_list, maxlen=max_len, dtype='float32', padding='post', value=0)

	# 转换标签为one-hot编码
	labels_one_hot = utils.to_categorical(labels, num_classes=num_classes)
	# print("labels_one_hot:")
	# print(labels_one_hot)
	# 计算分割点前,先确保数据集被完全构建
	num_elements = len(file_list_padded)
	train_size = int(num_elements * 0.8)
	
	# 先shuffle再分割
	indices = np.arange(num_elements)
	np.random.shuffle(indices)
	
	# 分割数据集
	train_indices = indices[:train_size]
	test_indices = indices[train_size:]
	
	# 获取训练和测试数据
	x_train = file_list_padded[train_indices]
	y_train = labels_one_hot[train_indices]
	x_test = file_list_padded[test_indices]
	y_test = labels_one_hot[test_indices]
	
	return x_train, y_train, x_test, y_test

def normalize_data(data, method='minmax', **kwargs):
	if method == 'minmax':
		scaler = MinMaxScaler(**kwargs)
	elif method == 'standard':
		scaler = StandardScaler(**kwargs)
	else:
		raise ValueError("Unsupported normalization method.")
	original_shape = data.shape
	data_reshaped = data.reshape(-1, data.shape[-1])
	normalized = scaler.fit_transform(data_reshaped)
	return normalized.reshape(original_shape)

def augment_data(data, noise_std=0.01, flip_prob=0.5):
	augmented = data.copy()
	noise = np.random.normal(0, noise_std, augmented.shape)
	augmented += noise
	if np.random.rand() < flip_prob:
		augmented[:, :, :3] *= -1
	return augmented

def create_dataset(data, labels, num_classes, batch_size=32, shuffle=True, augmentation_func=None):
	"""
	创建 tf.data.Dataset。labels 可以是整数数组,会被转换为 one-hot(需提供 num_classes)。
	augmentation_func 接受 numpy 数组并返回 numpy 数组(将通过 tf.py_function 调用)。
	"""
	if len(labels.shape) == 1:
		labels = tf.keras.utils.to_categorical(labels, num_classes)

	ds = tf.data.Dataset.from_tensor_slices((data.astype('float32'), labels.astype('float32')))
	if shuffle:
		ds = ds.shuffle(buffer_size=len(data))

	if augmentation_func:
		def _augment(x, y):
			x_aug = tf.py_function(lambda z: augmentation_func(z.numpy()) , [x], tf.float32)
			x_aug.set_shape(x.shape)
			return x_aug, y
		ds = ds.map(_augment, num_parallel_calls=tf.data.AUTOTUNE)

	ds = ds.batch(batch_size).prefetch(tf.data.AUTOTUNE)
	return ds

def preprocess_data(train_data_path, input_shape, num_classes,
					train_labels_path=None, val_data_path=None, val_labels_path=None,
					normalization_method='minmax', batch_size=32, augmentation=False, custom_augment_func=None):
	"""
	完整预处理管道,返回 train_ds, val_ds(val_ds 可能为 None)。
	"""
	train_data, train_labels = load_data(train_data_path, input_shape, train_labels_path)
	train_data = normalize_data(train_data, method=normalization_method)
	augment_func = custom_augment_func if custom_augment_func else (augment_data if augmentation else None)
	train_ds = create_dataset(train_data, train_labels, num_classes, batch_size=batch_size, shuffle=True, augmentation_func=augment_func)

	val_ds = None
	if val_data_path:
		val_data, val_labels = load_data(val_data_path, input_shape, val_labels_path)
		val_data = normalize_data(val_data, method=normalization_method)
		val_ds = create_dataset(val_data, val_labels, num_classes, batch_size=batch_size, shuffle=False, augmentation_func=None)

	return train_ds, val_ds
